# encoding: utf-8
# @File  : redis.py
# @Author: shaoyun
# @Date  :  2025/05/11
import asyncio
import json
import time
from typing import Any, Dict, List, Optional

import redis
from loguru import logger
from redis.asyncio import Redis

from redis.exceptions import ConnectionError, TimeoutError

from app.core.config import settings


class RedisClient:
    _client: Optional[Redis] = None

    @classmethod
    async def get_client(cls) -> Redis:
        """获取Redis客户端单例"""
        if cls._client is None:
            try:
                cls._client = Redis(
                    host=settings.REDIS_HOST,
                    port=settings.REDIS_PORT,
                    db=settings.REDIS_DB,
                    password=settings.REDIS_PASSWORD,
                    decode_responses=True
                )
                logger.info("Connected to Redis")
            except (ConnectionError, TimeoutError) as e:
                logger.error(f"Failed to connect to Redis: {e}")
                raise
        return cls._client

    @classmethod
    async def close(cls):
        """关闭Redis连接"""
        if cls._client:
            await cls._client.close()
            cls._client = None
            logger.info("Closed Redis connection")

    # ========== 任务队列操作 ==========

    @classmethod
    async def push_task(cls, task_id: str) -> int:
        """向任务队列推送任务ID"""
        client = await cls.get_client()
        return await client.rpush(settings.TASK_QUEUE_KEY, task_id)

    @classmethod
    async def pop_task(cls, timeout: int = 0) -> Optional[str]:
        """从任务队列获取任务ID（阻塞操作）"""
        client = await cls.get_client()
        if timeout > 0:
            result = await client.blpop(settings.TASK_QUEUE_KEY, timeout)
            return result[1] if result else None
        else:
            return await client.lpop(settings.TASK_QUEUE_KEY)

    @classmethod
    async def get_queue_size(cls) -> int:
        """获取任务队列大小"""
        try:
            client = await cls.get_client()
            normal_queue_size = await client.llen(settings.TASK_QUEUE_KEY)

            priority_queue_size = 0
            if hasattr(settings, 'PRIORITY_QUEUE_KEY'):
                priority_queue_size = await client.zcard(settings.PRIORITY_QUEUE_KEY)

            return normal_queue_size + priority_queue_size
        except Exception as e:
            logger.error(f"Error getting queue size: {e}")
            return 0

    # ========== 优先级任务队列 ==========

    @classmethod
    async def push_priority_task(cls, task_id: str, priority: int) -> bool:
        """将任务按优先级推入队列"""
        try:
            client = await cls.get_client()
            if hasattr(settings, 'PRIORITY_QUEUE_KEY'):
                await client.zadd(settings.PRIORITY_QUEUE_KEY, {task_id: priority})
                logger.debug(f"Task {task_id} pushed to priority queue with priority {priority}")
                return True
            else:
                await cls.push_task(task_id)
                return True
        except Exception as e:
            logger.error(f"Error pushing task to priority queue: {e}")
            return False

    @classmethod
    async def pop_priority_task(cls) -> Optional[str]:
        """从优先级队列中取出最高优先级的任务"""
        try:
            client = await cls.get_client()
            if hasattr(settings, 'PRIORITY_QUEUE_KEY'):
                result = await client.zpopmax(settings.PRIORITY_QUEUE_KEY)
                if result:
                    task_id, score = result[0]
                    logger.debug(f"Popped task {task_id} with priority {score}")
                    return task_id
            return None
        except Exception as e:
            logger.error(f"Error popping task from priority queue: {e}")
            return None

    # ========== 节点状态操作 ==========

    @classmethod
    async def set_node_info(cls, node_id: str, node_info: Dict[str, Any]) -> bool:
        """设置节点信息"""
        try:
            client = await cls.get_client()
            key = f"{settings.NODE_KEY_PREFIX}{node_id}"

            flat_info = {}
            for k, v in node_info.items():
                if isinstance(v, (dict, list)):
                    flat_info[k] = json.dumps(v)
                else:
                    flat_info[k] = str(v)

            await client.hmset(key, flat_info)
            await client.expire(key, settings.NODE_HEARTBEAT_TIMEOUT + 60)
            return True
        except Exception as e:
            logger.error(f"Error setting node info: {e}")
            return False

    @classmethod
    async def update_node_status(cls, node_id: str, status: Dict[str, Any]) -> bool:
        """更新节点状态"""
        client = await cls.get_client()
        key = f"{settings.NODE_KEY_PREFIX}{node_id}"

        for k, v in status.items():
            if isinstance(v, (dict, list)):
                status[k] = json.dumps(v)

        return await client.hmset(key, status)

    @classmethod
    async def get_node_status(cls, node_id: str) -> Dict[str, Any]:
        """获取节点状态"""
        client = await cls.get_client()
        key = f"{settings.NODE_KEY_PREFIX}{node_id}"
        result = await client.hgetall(key)

        for k, v in result.items():
            try:
                result[k] = json.loads(v)
            except (json.JSONDecodeError, TypeError):
                pass

        return result

    @classmethod
    async def get_all_nodes(cls) -> List[str]:
        """获取所有节点ID"""
        client = await cls.get_client()
        prefix = settings.NODE_KEY_PREFIX
        keys = await client.keys(f"{prefix}*")
        return [key[len(prefix):] for key in keys]

    # ========== 心跳操作 ==========

    @classmethod
    async def set_heartbeat(cls, node_id: str) -> bool:
        """设置节点心跳（简化版）"""
        try:
            client = await cls.get_client()
            key = f"{settings.HEARTBEAT_KEY_PREFIX}{node_id}"
            await client.setex(key, settings.NODE_HEARTBEAT_TIMEOUT, "1")
            return True
        except Exception as e:
            logger.error(f"Error setting heartbeat: {e}")
            return False

    @classmethod
    async def update_heartbeat(cls, node_id: str, timestamp: int) -> bool:
        """更新节点心跳"""
        client = await cls.get_client()
        key = f"{settings.HEARTBEAT_KEY_PREFIX}{node_id}"
        await client.set(key, timestamp)
        return await client.expire(key, settings.NODE_HEARTBEAT_TIMEOUT * 2)

    @classmethod
    async def get_heartbeat(cls, node_id: str) -> Optional[int]:
        """获取节点最后心跳时间"""
        client = await cls.get_client()
        key = f"{settings.HEARTBEAT_KEY_PREFIX}{node_id}"
        value = await client.get(key)
        return int(value) if value else None

    @classmethod
    async def check_heartbeat(cls, node_id: str) -> bool:
        """检查节点心跳是否存在"""
        try:
            client = await cls.get_client()
            key = f"{settings.HEARTBEAT_KEY_PREFIX}{node_id}"
            return await client.exists(key) > 0
        except Exception as e:
            logger.error(f"Error checking heartbeat: {e}")
            return False

    # ========== 缓存操作 ==========

    @classmethod
    async def set_cache(cls, key: str, value: Any, expire: int = 3600) -> bool:
        """设置缓存"""
        try:
            client = await cls.get_client()
            cache_key = f"cache:{key}"
            if isinstance(value, (dict, list)):
                value = json.dumps(value)
            return await client.setex(cache_key, expire, value)
        except Exception as e:
            logger.error(f"Error setting cache: {e}")
            return False

    @classmethod
    async def get_cache(cls, key: str) -> Optional[Any]:
        """获取缓存"""
        try:
            client = await cls.get_client()
            cache_key = f"cache:{key}"
            value = await client.get(cache_key)
            if value:
                try:
                    return json.loads(value)
                except (json.JSONDecodeError, TypeError):
                    return value
            return None
        except Exception as e:
            logger.error(f"Error getting cache: {e}")
            return None

    # ========== 通用操作 ==========

    @classmethod
    async def set_key(cls, key: str, value: Any, ex: Optional[int] = None) -> bool:
        """设置键值对"""
        client = await cls.get_client()
        if isinstance(value, (dict, list)):
            value = json.dumps(value)
        return await client.set(key, value, ex=ex)

    @classmethod
    async def get_key(cls, key: str) -> Optional[str]:
        """获取键值"""
        client = await cls.get_client()
        return await client.get(key)

    @classmethod
    async def delete_key(cls, key: str) -> int:
        """删除键"""
        client = await cls.get_client()
        return await client.delete(key)

    # ========== 发布/订阅操作 ==========

    @classmethod
    async def publish_event(cls, channel: str, data: Dict[str, Any]) -> int:
        """发布事件到指定频道"""
        try:
            client = await cls.get_client()
            message = json.dumps(data)
            subscriber_count = await client.publish(channel, message)
            logger.debug(f"Published event to channel '{channel}' with data: {data}, subscribers: {subscriber_count}")
            return subscriber_count
        except Exception as e:
            logger.error(f"Error publishing event to channel '{channel}': {e}")
            return 0

    @classmethod
    async def get_pubsub(cls) -> redis.client.PubSub:
        """获取发布/订阅对象"""
        try:
            client = await cls.get_client()
            pubsub = client.pubsub()
            return pubsub
        except Exception as e:
            logger.error(f"Error creating pubsub object: {e}")
            raise

    @classmethod
    async def subscribe_to_channels(cls, *channels: str) -> redis.client.PubSub:
        """订阅一个或多个频道"""
        try:
            pubsub = await cls.get_pubsub()
            await pubsub.subscribe(*channels)
            logger.info(f"Subscribed to channels: {', '.join(channels)}")
            return pubsub
        except Exception as e:
            logger.error(f"Error subscribing to channels: {e}")
            raise

    @classmethod
    def get_message(cls, pubsub: redis.client.PubSub, timeout: float = 0.0) -> Optional[Dict[str, Any]]:
        """从订阅的频道获取消息"""
        try:
            message = pubsub.get_message(timeout=timeout)
            if message and message['type'] == 'message':
                try:
                    message['data'] = json.loads(message['data'])
                except (json.JSONDecodeError, TypeError):
                    pass
            return message
        except Exception as e:
            logger.error(f"Error getting message from pubsub: {e}")
            return None

    @classmethod
    async def publish_spider_sync_event(cls, spider_id: str, file_id: str, action: str = "sync") -> int:
        """发布爬虫同步事件（专用方法）"""
        event_data = {
            "spider_id": spider_id,
            "file_id": file_id,
            "action": action,
            "timestamp": int(time.time())
        }
        return await cls.publish_event("spider.sync", event_data)

    @classmethod
    async def publish_file_upload_event(cls, spider_id: str, file_id: str) -> int:
        """发布文件上传事件（专用方法）"""
        event_data = {
            "spider_id": spider_id,
            "file_id": file_id,
            "action": "upload",
            "timestamp": int(time.time())
        }
        return await cls.publish_event("file.upload", event_data)

class AsyncRedisPubSub:
    """异步Redis发布/订阅助手类"""

    def __init__(self, channels: List[str]):
        self.channels = channels
        self.pubsub = None
        self._running = False

    async def start(self):
        """开始订阅"""
        self.pubsub = await RedisClient.get_pubsub()
        await self.pubsub.subscribe(*self.channels)
        self._running = True
        logger.info(f"Started async subscription to channels: {', '.join(self.channels)}")

    async def stop(self):
        """停止订阅"""
        self._running = False
        if self.pubsub:
            await self.pubsub.unsubscribe()
            await self.pubsub.close()
            logger.info("Stopped async subscription")

    async def listen(self):
        """异步监听消息"""
        while self._running:
            try:
                message = await self.pubsub.get_message(timeout=0.1)
                if message and message['type'] == 'message':
                    try:
                        data = json.loads(message['data'])
                        yield {
                            'channel': message['channel'].decode('utf-8'),
                            'data': data
                        }
                    except (json.JSONDecodeError, TypeError, AttributeError):
                        logger.warning(f"Failed to decode message: {message}")
                await asyncio.sleep(0.01)
            except Exception as e:
                logger.error(f"Error in async listen: {e}")
                await asyncio.sleep(1)